package com.wyl.mybatis.intercept;

import cn.hutool.db.sql.SqlBuilder;
import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
import com.wyl.mybatis.config.SqlSessionFactoryConfig;
import com.wyl.mybatis.page.Page;
import com.wyl.mybatis.util.SqlParamUtil;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.util.List;
import java.util.Map;

/**
 * @Description
 * @Author WuYiLong
 * @Date 2024/3/13 12:06
 */
@Intercepts({
        @Signature(
                type = Executor.class,
                method = "query",
                args = {
                        MappedStatement.class,
                        Object.class,
                        RowBounds.class,
                        ResultHandler.class
                }),
        @Signature(
                type = Executor.class,
                method = "query",
                args = {MappedStatement.class,
                        Object.class,
                        RowBounds.class,
                        ResultHandler.class,
                        CacheKey.class,
                        BoundSql.class})})
public class PageIntercept implements Interceptor {

    private final static Logger log = LoggerFactory.getLogger(PageIntercept.class);

    @Override
    public Object intercept(Invocation invocation) throws Throwable {

        Executor executor = (Executor) invocation.getTarget();
        Object[] args = invocation.getArgs();
        MappedStatement ms = (MappedStatement) args[0];
        Object params = args[1];
        RowBounds rowBounds = (RowBounds) args[2];
        BoundSql boundSql = ms.getBoundSql(params);
        String sql = boundSql.getSql();
        ResultHandler resultHandler = (ResultHandler) args[3];

        Page page = null;
        if(params != null) {
            Map<String, Object> paramMap = SqlParamUtil.filter(params);
            for (Map.Entry<String, Object> mapEntry : paramMap.entrySet()) {
                Object v = mapEntry.getValue();
                if (v instanceof Page) {
                    page = (Page) v;

                }
            }
        }

        if(page != null) {
            String countSql = countSql(sql);
            int count = 0;
            SqlSessionFactory sqlSessionFactory = SqlSessionFactoryConfig.buildSqlSessionFactory();
            try (SqlSession sqlSession = sqlSessionFactory.openSession()) {
                Connection connection = sqlSession.getConnection();
                PreparedStatement preparedStatement = connection.prepareStatement(countSql);
                ResultSet resultSet = preparedStatement.executeQuery();
                while (resultSet.next()) {
                    count = resultSet.getInt("count");
                }
            }

            page.setTotal(count);
            rowBounds = new RowBounds(page.getCurrentPage()-1,page.getPageSize());
        }

        CacheKey cacheKey;
        if (args.length == 4) {
            cacheKey = executor.createCacheKey(ms, params, rowBounds, boundSql);
        } else {
            cacheKey = (CacheKey) args[4];
            boundSql = (BoundSql) args[5];
        }

        List<Object> query = executor.query(ms, params, rowBounds, resultHandler, cacheKey, boundSql);
        return query;
    }

    @Override
    public Object plugin(Object target) {
        if (target instanceof Executor) {
            return Plugin.wrap(target, this);
        }
        return target;
    }

    private String countSql(String sql) {
        SqlBuilder sqlBuilder = new SqlBuilder();
        sqlBuilder.select("count(*) count");
        List<SQLStatement> sqlStatements = SQLUtils.parseStatements(sql, DbType.mysql);
        SQLSelectStatement sqlStatement = (SQLSelectStatement) sqlStatements.get(0);
        MySqlSelectQueryBlock queryBlock = (MySqlSelectQueryBlock) sqlStatement.getSelect().getQueryBlock();
        sqlBuilder.from(queryBlock.getFrom().toString());
        if (queryBlock.getWhere() != null) {
            sqlBuilder.where(queryBlock.getWhere().toString());
        }
        return sqlBuilder.build();
    }
}
